home *** CD-ROM | disk | FTP | other *** search
/ Chip 2007 January, February, March & April / Chip-Cover-CD-2007-02.iso / Pakiet bezpieczenstwa / mini Pentoo LiveCD 2006.1 / mpentoo-2006.1.iso / livecd.squashfs / usr / lib / python2.4 / test / test_itertools.py < prev    next >
Encoding:
Python Source  |  2005-10-18  |  32.2 KB  |  942 lines
  1. import unittest
  2. from test import test_support
  3. from itertools import *
  4. from weakref import proxy
  5. import sys
  6. import operator
  7. import random
  8.  
  9. def onearg(x):
  10.     'Test function of one argument'
  11.     return 2*x
  12.  
  13. def errfunc(*args):
  14.     'Test function that raises an error'
  15.     raise ValueError
  16.  
  17. def gen3():
  18.     'Non-restartable source sequence'
  19.     for i in (0, 1, 2):
  20.         yield i
  21.  
  22. def isEven(x):
  23.     'Test predicate'
  24.     return x%2==0
  25.  
  26. def isOdd(x):
  27.     'Test predicate'
  28.     return x%2==1
  29.  
  30. class StopNow:
  31.     'Class emulating an empty iterable.'
  32.     def __iter__(self):
  33.         return self
  34.     def next(self):
  35.         raise StopIteration
  36.  
  37. def take(n, seq):
  38.     'Convenience function for partially consuming a long of infinite iterable'
  39.     return list(islice(seq, n))
  40.  
  41. class TestBasicOps(unittest.TestCase):
  42.     def test_chain(self):
  43.         self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
  44.         self.assertEqual(list(chain('abc')), list('abc'))
  45.         self.assertEqual(list(chain('')), [])
  46.         self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
  47.         self.assertRaises(TypeError, chain, 2, 3)
  48.  
  49.     def test_count(self):
  50.         self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
  51.         self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
  52.         self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
  53.         self.assertRaises(TypeError, count, 2, 3)
  54.         self.assertRaises(TypeError, count, 'a')
  55.         c = count(sys.maxint-2)   # verify that rollover doesn't crash
  56.         c.next(); c.next(); c.next(); c.next(); c.next()
  57.         c = count(3)
  58.         self.assertEqual(repr(c), 'count(3)')
  59.         c.next()
  60.         self.assertEqual(repr(c), 'count(4)')
  61.  
  62.     def test_cycle(self):
  63.         self.assertEqual(take(10, cycle('abc')), list('abcabcabca'))
  64.         self.assertEqual(list(cycle('')), [])
  65.         self.assertRaises(TypeError, cycle)
  66.         self.assertRaises(TypeError, cycle, 5)
  67.         self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0])
  68.  
  69.     def test_groupby(self):
  70.         # Check whether it accepts arguments correctly
  71.         self.assertEqual([], list(groupby([])))
  72.         self.assertEqual([], list(groupby([], key=id)))
  73.         self.assertRaises(TypeError, list, groupby('abc', []))
  74.         self.assertRaises(TypeError, groupby, None)
  75.         self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
  76.  
  77.         # Check normal input
  78.         s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
  79.              (2,15,22), (3,16,23), (3,17,23)]
  80.         dup = []
  81.         for k, g in groupby(s, lambda r:r[0]):
  82.             for elem in g:
  83.                 self.assertEqual(k, elem[0])
  84.                 dup.append(elem)
  85.         self.assertEqual(s, dup)
  86.  
  87.         # Check nested case
  88.         dup = []
  89.         for k, g in groupby(s, lambda r:r[0]):
  90.             for ik, ig in groupby(g, lambda r:r[2]):
  91.                 for elem in ig:
  92.                     self.assertEqual(k, elem[0])
  93.                     self.assertEqual(ik, elem[2])
  94.                     dup.append(elem)
  95.         self.assertEqual(s, dup)
  96.  
  97.         # Check case where inner iterator is not used
  98.         keys = [k for k, g in groupby(s, lambda r:r[0])]
  99.         expectedkeys = set([r[0] for r in s])
  100.         self.assertEqual(set(keys), expectedkeys)
  101.         self.assertEqual(len(keys), len(expectedkeys))
  102.  
  103.         # Exercise pipes and filters style
  104.         s = 'abracadabra'
  105.         # sort s | uniq
  106.         r = [k for k, g in groupby(sorted(s))]
  107.         self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
  108.         # sort s | uniq -d
  109.         r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
  110.         self.assertEqual(r, ['a', 'b', 'r'])
  111.         # sort s | uniq -c
  112.         r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
  113.         self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
  114.         # sort s | uniq -c | sort -rn | head -3
  115.         r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
  116.         self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
  117.  
  118.         # iter.next failure
  119.         class ExpectedError(Exception):
  120.             pass
  121.         def delayed_raise(n=0):
  122.             for i in range(n):
  123.                 yield 'yo'
  124.             raise ExpectedError
  125.         def gulp(iterable, keyp=None, func=list):
  126.             return [func(g) for k, g in groupby(iterable, keyp)]
  127.  
  128.         # iter.next failure on outer object
  129.         self.assertRaises(ExpectedError, gulp, delayed_raise(0))
  130.         # iter.next failure on inner object
  131.         self.assertRaises(ExpectedError, gulp, delayed_raise(1))
  132.  
  133.         # __cmp__ failure
  134.         class DummyCmp:
  135.             def __cmp__(self, dst):
  136.                 raise ExpectedError
  137.         s = [DummyCmp(), DummyCmp(), None]
  138.  
  139.         # __cmp__ failure on outer object
  140.         self.assertRaises(ExpectedError, gulp, s, func=id)
  141.         # __cmp__ failure on inner object
  142.         self.assertRaises(ExpectedError, gulp, s)
  143.  
  144.         # keyfunc failure
  145.         def keyfunc(obj):
  146.             if keyfunc.skip > 0:
  147.                 keyfunc.skip -= 1
  148.                 return obj
  149.             else:
  150.                 raise ExpectedError
  151.  
  152.         # keyfunc failure on outer object
  153.         keyfunc.skip = 0
  154.         self.assertRaises(ExpectedError, gulp, [None], keyfunc)
  155.         keyfunc.skip = 1
  156.         self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
  157.  
  158.     def test_ifilter(self):
  159.         self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
  160.         self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
  161.         self.assertEqual(take(4, ifilter(isEven, count())), [0,2,4,6])
  162.         self.assertRaises(TypeError, ifilter)
  163.         self.assertRaises(TypeError, ifilter, lambda x:x)
  164.         self.assertRaises(TypeError, ifilter, lambda x:x, range(6), 7)
  165.         self.assertRaises(TypeError, ifilter, isEven, 3)
  166.         self.assertRaises(TypeError, ifilter(range(6), range(6)).next)
  167.  
  168.     def test_ifilterfalse(self):
  169.         self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
  170.         self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
  171.         self.assertEqual(take(4, ifilterfalse(isEven, count())), [1,3,5,7])
  172.         self.assertRaises(TypeError, ifilterfalse)
  173.         self.assertRaises(TypeError, ifilterfalse, lambda x:x)
  174.         self.assertRaises(TypeError, ifilterfalse, lambda x:x, range(6), 7)
  175.         self.assertRaises(TypeError, ifilterfalse, isEven, 3)
  176.         self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
  177.  
  178.     def test_izip(self):
  179.         ans = [(x,y) for x, y in izip('abc',count())]
  180.         self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
  181.         self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
  182.         self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
  183.         self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
  184.         self.assertEqual(list(izip('abcdef')), zip('abcdef'))
  185.         self.assertEqual(list(izip()), zip())
  186.         self.assertRaises(TypeError, izip, 3)
  187.         self.assertRaises(TypeError, izip, range(3), 3)
  188.         # Check tuple re-use (implementation detail)
  189.         self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
  190.                          zip('abc', 'def'))
  191.         self.assertEqual([pair for pair in izip('abc', 'def')],
  192.                          zip('abc', 'def'))
  193.         ids = map(id, izip('abc', 'def'))
  194.         self.assertEqual(min(ids), max(ids))
  195.         ids = map(id, list(izip('abc', 'def')))
  196.         self.assertEqual(len(dict.fromkeys(ids)), len(ids))
  197.  
  198.     def test_repeat(self):
  199.         self.assertEqual(zip(xrange(3),repeat('a')),
  200.                          [(0, 'a'), (1, 'a'), (2, 'a')])
  201.         self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
  202.         self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
  203.         self.assertEqual(list(repeat('a', 0)), [])
  204.         self.assertEqual(list(repeat('a', -3)), [])
  205.         self.assertRaises(TypeError, repeat)
  206.         self.assertRaises(TypeError, repeat, None, 3, 4)
  207.         self.assertRaises(TypeError, repeat, None, 'a')
  208.         r = repeat(1+0j)
  209.         self.assertEqual(repr(r), 'repeat((1+0j))')
  210.         r = repeat(1+0j, 5)
  211.         self.assertEqual(repr(r), 'repeat((1+0j), 5)')
  212.         list(r)
  213.         self.assertEqual(repr(r), 'repeat((1+0j), 0)')
  214.  
  215.     def test_imap(self):
  216.         self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
  217.                          [0**1, 1**2, 2**3])
  218.         self.assertEqual(list(imap(None, 'abc', range(5))),
  219.                          [('a',0),('b',1),('c',2)])
  220.         self.assertEqual(list(imap(None, 'abc', count())),
  221.                          [('a',0),('b',1),('c',2)])
  222.         self.assertEqual(take(2,imap(None, 'abc', count())),
  223.                          [('a',0),('b',1)])
  224.         self.assertEqual(list(imap(operator.pow, [])), [])
  225.         self.assertRaises(TypeError, imap)
  226.         self.assertRaises(TypeError, imap, operator.neg)
  227.         self.assertRaises(TypeError, imap(10, range(5)).next)
  228.         self.assertRaises(ValueError, imap(errfunc, [4], [5]).next)
  229.         self.assertRaises(TypeError, imap(onearg, [4], [5]).next)
  230.  
  231.     def test_starmap(self):
  232.         self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
  233.                          [0**1, 1**2, 2**3])
  234.         self.assertEqual(take(3, starmap(operator.pow, izip(count(), count(1)))),
  235.                          [0**1, 1**2, 2**3])
  236.         self.assertEqual(list(starmap(operator.pow, [])), [])
  237.         self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
  238.         self.assertRaises(TypeError, starmap)
  239.         self.assertRaises(TypeError, starmap, operator.pow, [(4,5)], 'extra')
  240.         self.assertRaises(TypeError, starmap(10, [(4,5)]).next)
  241.         self.assertRaises(ValueError, starmap(errfunc, [(4,5)]).next)
  242.         self.assertRaises(TypeError, starmap(onearg, [(4,5)]).next)
  243.  
  244.     def test_islice(self):
  245.         for args in [          # islice(args) should agree with range(args)
  246.                 (10, 20, 3),
  247.                 (10, 3, 20),
  248.                 (10, 20),
  249.                 (10, 3),
  250.                 (20,)
  251.                 ]:
  252.             self.assertEqual(list(islice(xrange(100), *args)), range(*args))
  253.  
  254.         for args, tgtargs in [  # Stop when seqn is exhausted
  255.                 ((10, 110, 3), ((10, 100, 3))),
  256.                 ((10, 110), ((10, 100))),
  257.                 ((110,), (100,))
  258.                 ]:
  259.             self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
  260.  
  261.         # Test stop=None
  262.         self.assertEqual(list(islice(xrange(10), None)), range(10))
  263.         self.assertEqual(list(islice(xrange(10), 2, None)), range(2, 10))
  264.         self.assertEqual(list(islice(xrange(10), 1, None, 2)), range(1, 10, 2))
  265.  
  266.         # Test invalid arguments
  267.         self.assertRaises(TypeError, islice, xrange(10))
  268.         self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
  269.         self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
  270.         self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
  271.         self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
  272.         self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
  273.         self.assertRaises(ValueError, islice, xrange(10), 'a')
  274.         self.assertRaises(ValueError, islice, xrange(10), 'a', 1)
  275.         self.assertRaises(ValueError, islice, xrange(10), 1, 'a')
  276.         self.assertRaises(ValueError, islice, xrange(10), 'a', 1, 1)
  277.         self.assertRaises(ValueError, islice, xrange(10), 1, 'a', 1)
  278.         self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
  279.  
  280.     def test_takewhile(self):
  281.         data = [1, 3, 5, 20, 2, 4, 6, 8]
  282.         underten = lambda x: x<10
  283.         self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
  284.         self.assertEqual(list(takewhile(underten, [])), [])
  285.         self.assertRaises(TypeError, takewhile)
  286.         self.assertRaises(TypeError, takewhile, operator.pow)
  287.         self.assertRaises(TypeError, takewhile, operator.pow, [(4,5)], 'extra')
  288.         self.assertRaises(TypeError, takewhile(10, [(4,5)]).next)
  289.         self.assertRaises(ValueError, takewhile(errfunc, [(4,5)]).next)
  290.         t = takewhile(bool, [1, 1, 1, 0, 0, 0])
  291.         self.assertEqual(list(t), [1, 1, 1])
  292.         self.assertRaises(StopIteration, t.next)
  293.  
  294.     def test_dropwhile(self):
  295.         data = [1, 3, 5, 20, 2, 4, 6, 8]
  296.         underten = lambda x: x<10
  297.         self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
  298.         self.assertEqual(list(dropwhile(underten, [])), [])
  299.         self.assertRaises(TypeError, dropwhile)
  300.         self.assertRaises(TypeError, dropwhile, operator.pow)
  301.         self.assertRaises(TypeError, dropwhile, operator.pow, [(4,5)], 'extra')
  302.         self.assertRaises(TypeError, dropwhile(10, [(4,5)]).next)
  303.         self.assertRaises(ValueError, dropwhile(errfunc, [(4,5)]).next)
  304.  
  305.     def test_tee(self):
  306.         n = 200
  307.         def irange(n):
  308.             for i in xrange(n):
  309.                 yield i
  310.  
  311.         a, b = tee([])        # test empty iterator
  312.         self.assertEqual(list(a), [])
  313.         self.assertEqual(list(b), [])
  314.  
  315.         a, b = tee(irange(n)) # test 100% interleaved
  316.         self.assertEqual(zip(a,b), zip(range(n),range(n)))
  317.  
  318.         a, b = tee(irange(n)) # test 0% interleaved
  319.         self.assertEqual(list(a), range(n))
  320.         self.assertEqual(list(b), range(n))
  321.  
  322.         a, b = tee(irange(n)) # test dealloc of leading iterator
  323.         for i in xrange(100):
  324.             self.assertEqual(a.next(), i)
  325.         del a
  326.         self.assertEqual(list(b), range(n))
  327.  
  328.         a, b = tee(irange(n)) # test dealloc of trailing iterator
  329.         for i in xrange(100):
  330.             self.assertEqual(a.next(), i)
  331.         del b
  332.         self.assertEqual(list(a), range(100, n))
  333.  
  334.         for j in xrange(5):   # test randomly interleaved
  335.             order = [0]*n + [1]*n
  336.             random.shuffle(order)
  337.             lists = ([], [])
  338.             its = tee(irange(n))
  339.             for i in order:
  340.                 value = its[i].next()
  341.                 lists[i].append(value)
  342.             self.assertEqual(lists[0], range(n))
  343.             self.assertEqual(lists[1], range(n))
  344.  
  345.         # test argument format checking
  346.         self.assertRaises(TypeError, tee)
  347.         self.assertRaises(TypeError, tee, 3)
  348.         self.assertRaises(TypeError, tee, [1,2], 'x')
  349.         self.assertRaises(TypeError, tee, [1,2], 3, 'x')
  350.  
  351.         # tee object should be instantiable
  352.         a, b = tee('abc')
  353.         c = type(a)('def')
  354.         self.assertEqual(list(c), list('def'))
  355.  
  356.         # test long-lagged and multi-way split
  357.         a, b, c = tee(xrange(2000), 3)
  358.         for i in xrange(100):
  359.             self.assertEqual(a.next(), i)
  360.         self.assertEqual(list(b), range(2000))
  361.         self.assertEqual([c.next(), c.next()], range(2))
  362.         self.assertEqual(list(a), range(100,2000))
  363.         self.assertEqual(list(c), range(2,2000))
  364.  
  365.         # test values of n
  366.         self.assertRaises(TypeError, tee, 'abc', 'invalid')
  367.         for n in xrange(5):
  368.             result = tee('abc', n)
  369.             self.assertEqual(type(result), tuple)
  370.             self.assertEqual(len(result), n)
  371.             self.assertEqual(map(list, result), [list('abc')]*n)
  372.  
  373.         # tee pass-through to copyable iterator
  374.         a, b = tee('abc')
  375.         c, d = tee(a)
  376.         self.assert_(a is c)
  377.  
  378.         # test tee_new
  379.         t1, t2 = tee('abc')
  380.         tnew = type(t1)
  381.         self.assertRaises(TypeError, tnew)
  382.         self.assertRaises(TypeError, tnew, 10)
  383.         t3 = tnew(t1)
  384.         self.assert_(list(t1) == list(t2) == list(t3) == list('abc'))
  385.  
  386.         # test that tee objects are weak referencable
  387.         a, b = tee(xrange(10))
  388.         p = proxy(a)
  389.         self.assertEqual(getattr(p, '__class__'), type(b))
  390.         del a
  391.         self.assertRaises(ReferenceError, getattr, p, '__class__')
  392.  
  393.     def test_StopIteration(self):
  394.         self.assertRaises(StopIteration, izip().next)
  395.  
  396.         for f in (chain, cycle, izip, groupby):
  397.             self.assertRaises(StopIteration, f([]).next)
  398.             self.assertRaises(StopIteration, f(StopNow()).next)
  399.  
  400.         self.assertRaises(StopIteration, islice([], None).next)
  401.         self.assertRaises(StopIteration, islice(StopNow(), None).next)
  402.  
  403.         p, q = tee([])
  404.         self.assertRaises(StopIteration, p.next)
  405.         self.assertRaises(StopIteration, q.next)
  406.         p, q = tee(StopNow())
  407.         self.assertRaises(StopIteration, p.next)
  408.         self.assertRaises(StopIteration, q.next)
  409.  
  410.         self.assertRaises(StopIteration, repeat(None, 0).next)
  411.  
  412.         for f in (ifilter, ifilterfalse, imap, takewhile, dropwhile, starmap):
  413.             self.assertRaises(StopIteration, f(lambda x:x, []).next)
  414.             self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next)
  415.  
  416. class TestGC(unittest.TestCase):
  417.  
  418.     def makecycle(self, iterator, container):
  419.         container.append(iterator)
  420.         iterator.next()
  421.         del container, iterator
  422.  
  423.     def test_chain(self):
  424.         a = []
  425.         self.makecycle(chain(a), a)
  426.  
  427.     def test_cycle(self):
  428.         a = []
  429.         self.makecycle(cycle([a]*2), a)
  430.  
  431.     def test_dropwhile(self):
  432.         a = []
  433.         self.makecycle(dropwhile(bool, [0, a, a]), a)
  434.  
  435.     def test_groupby(self):
  436.         a = []
  437.         self.makecycle(groupby([a]*2, lambda x:x), a)
  438.  
  439.     def test_ifilter(self):
  440.         a = []
  441.         self.makecycle(ifilter(lambda x:True, [a]*2), a)
  442.  
  443.     def test_ifilterfalse(self):
  444.         a = []
  445.         self.makecycle(ifilterfalse(lambda x:False, a), a)
  446.  
  447.     def test_izip(self):
  448.         a = []
  449.         self.makecycle(izip([a]*2, [a]*3), a)
  450.  
  451.     def test_imap(self):
  452.         a = []
  453.         self.makecycle(imap(lambda x:x, [a]*2), a)
  454.  
  455.     def test_islice(self):
  456.         a = []
  457.         self.makecycle(islice([a]*2, None), a)
  458.  
  459.     def test_repeat(self):
  460.         a = []
  461.         self.makecycle(repeat(a), a)
  462.  
  463.     def test_starmap(self):
  464.         a = []
  465.         self.makecycle(starmap(lambda *t: t, [(a,a)]*2), a)
  466.  
  467.     def test_takewhile(self):
  468.         a = []
  469.         self.makecycle(takewhile(bool, [1, 0, a, a]), a)
  470.  
  471. def R(seqn):
  472.     'Regular generator'
  473.     for i in seqn:
  474.         yield i
  475.  
  476. class G:
  477.     'Sequence using __getitem__'
  478.     def __init__(self, seqn):
  479.         self.seqn = seqn
  480.     def __getitem__(self, i):
  481.         return self.seqn[i]
  482.  
  483. class I:
  484.     'Sequence using iterator protocol'
  485.     def __init__(self, seqn):
  486.         self.seqn = seqn
  487.         self.i = 0
  488.     def __iter__(self):
  489.         return self
  490.     def next(self):
  491.         if self.i >= len(self.seqn): raise StopIteration
  492.         v = self.seqn[self.i]
  493.         self.i += 1
  494.         return v
  495.  
  496. class Ig:
  497.     'Sequence using iterator protocol defined with a generator'
  498.     def __init__(self, seqn):
  499.         self.seqn = seqn
  500.         self.i = 0
  501.     def __iter__(self):
  502.         for val in self.seqn:
  503.             yield val
  504.  
  505. class X:
  506.     'Missing __getitem__ and __iter__'
  507.     def __init__(self, seqn):
  508.         self.seqn = seqn
  509.         self.i = 0
  510.     def next(self):
  511.         if self.i >= len(self.seqn): raise StopIteration
  512.         v = self.seqn[self.i]
  513.         self.i += 1
  514.         return v
  515.  
  516. class N:
  517.     'Iterator missing next()'
  518.     def __init__(self, seqn):
  519.         self.seqn = seqn
  520.         self.i = 0
  521.     def __iter__(self):
  522.         return self
  523.  
  524. class E:
  525.     'Test propagation of exceptions'
  526.     def __init__(self, seqn):
  527.         self.seqn = seqn
  528.         self.i = 0
  529.     def __iter__(self):
  530.         return self
  531.     def next(self):
  532.         3 // 0
  533.  
  534. class S:
  535.     'Test immediate stop'
  536.     def __init__(self, seqn):
  537.         pass
  538.     def __iter__(self):
  539.         return self
  540.     def next(self):
  541.         raise StopIteration
  542.  
  543. def L(seqn):
  544.     'Test multiple tiers of iterators'
  545.     return chain(imap(lambda x:x, R(Ig(G(seqn)))))
  546.  
  547.  
  548. class TestVariousIteratorArgs(unittest.TestCase):
  549.  
  550.     def test_chain(self):
  551.         for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
  552.             for g in (G, I, Ig, S, L, R):
  553.                 self.assertEqual(list(chain(g(s))), list(g(s)))
  554.                 self.assertEqual(list(chain(g(s), g(s))), list(g(s))+list(g(s)))
  555.             self.assertRaises(TypeError, chain, X(s))
  556.             self.assertRaises(TypeError, list, chain(N(s)))
  557.             self.assertRaises(ZeroDivisionError, list, chain(E(s)))
  558.  
  559.     def test_cycle(self):
  560.         for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
  561.             for g in (G, I, Ig, S, L, R):
  562.                 tgtlen = len(s) * 3
  563.                 expected = list(g(s))*3
  564.                 actual = list(islice(cycle(g(s)), tgtlen))
  565.                 self.assertEqual(actual, expected)
  566.             self.assertRaises(TypeError, cycle, X(s))
  567.             self.assertRaises(TypeError, list, cycle(N(s)))
  568.             self.assertRaises(ZeroDivisionError, list, cycle(E(s)))
  569.  
  570.     def test_groupby(self):
  571.         for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
  572.             for g in (G, I, Ig, S, L, R):
  573.                 self.assertEqual([k for k, sb in groupby(g(s))], list(g(s)))
  574.             self.assertRaises(TypeError, groupby, X(s))
  575.             self.assertRaises(TypeError, list, groupby(N(s)))
  576.             self.assertRaises(ZeroDivisionError, list, groupby(E(s)))
  577.  
  578.     def test_ifilter(self):
  579.         for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
  580.             for g in (G, I, Ig, S, L, R):
  581.                 self.assertEqual(list(ifilter(isEven, g(s))), filter(isEven, g(s)))
  582.             self.assertRaises(TypeError, ifilter, isEven, X(s))
  583.             self.assertRaises(TypeError, list, ifilter(isEven, N(s)))
  584.             self.assertRaises(ZeroDivisionError, list, ifilter(isEven, E(s)))
  585.  
  586.     def test_ifilterfalse(self):
  587.         for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
  588.             for g in (G, I, Ig, S, L, R):
  589.                 self.assertEqual(list(ifilterfalse(isEven, g(s))), filter(isOdd, g(s)))
  590.             self.assertRaises(TypeError, ifilterfalse, isEven, X(s))
  591.             self.assertRaises(TypeError, list, ifilterfalse(isEven, N(s)))
  592.             self.assertRaises(ZeroDivisionError, list, ifilterfalse(isEven, E(s)))
  593.  
  594.     def test_izip(self):
  595.         for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
  596.             for g in (G, I, Ig, S, L, R):
  597.                 self.assertEqual(list(izip(g(s))), zip(g(s)))
  598.                 self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
  599.             self.assertRaises(TypeError, izip, X(s))
  600.             self.assertRaises(TypeError, list, izip(N(s)))
  601.             self.assertRaises(ZeroDivisionError, list, izip(E(s)))
  602.  
  603.     def test_imap(self):
  604.         for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
  605.             for g in (G, I, Ig, S, L, R):
  606.                 self.assertEqual(list(imap(onearg, g(s))), map(onearg, g(s)))
  607.                 self.assertEqual(list(imap(operator.pow, g(s), g(s))), map(operator.pow, g(s), g(s)))
  608.             self.assertRaises(TypeError, imap, onearg, X(s))
  609.             self.assertRaises(TypeError, list, imap(onearg, N(s)))
  610.             self.assertRaises(ZeroDivisionError, list, imap(onearg, E(s)))
  611.  
  612.     def test_islice(self):
  613.         for s in ("12345", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
  614.             for g in (G, I, Ig, S, L, R):
  615.                 self.assertEqual(list(islice(g(s),1,None,2)), list(g(s))[1::2])
  616.             self.assertRaises(TypeError, islice, X(s), 10)
  617.             self.assertRaises(TypeError, list, islice(N(s), 10))
  618.             self.assertRaises(ZeroDivisionError, list, islice(E(s), 10))
  619.  
  620.     def test_starmap(self):
  621.         for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
  622.             for g in (G, I, Ig, S, L, R):
  623.                 ss = zip(s, s)
  624.                 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
  625.             self.assertRaises(TypeError, starmap, operator.pow, X(ss))
  626.             self.assertRaises(TypeError, list, starmap(operator.pow, N(ss)))
  627.             self.assertRaises(ZeroDivisionError, list, starmap(operator.pow, E(ss)))
  628.  
  629.     def test_takewhile(self):
  630.         for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
  631.             for g in (G, I, Ig, S, L, R):
  632.                 tgt = []
  633.                 for elem in g(s):
  634.                     if not isEven(elem): break
  635.                     tgt.append(elem)
  636.                 self.assertEqual(list(takewhile(isEven, g(s))), tgt)
  637.             self.assertRaises(TypeError, takewhile, isEven, X(s))
  638.             self.assertRaises(TypeError, list, takewhile(isEven, N(s)))
  639.             self.assertRaises(ZeroDivisionError, list, takewhile(isEven, E(s)))
  640.  
  641.     def test_dropwhile(self):
  642.         for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)):
  643.             for g in (G, I, Ig, S, L, R):
  644.                 tgt = []
  645.                 for elem in g(s):
  646.                     if not tgt and isOdd(elem): continue
  647.                     tgt.append(elem)
  648.                 self.assertEqual(list(dropwhile(isOdd, g(s))), tgt)
  649.             self.assertRaises(TypeError, dropwhile, isOdd, X(s))
  650.             self.assertRaises(TypeError, list, dropwhile(isOdd, N(s)))
  651.             self.assertRaises(ZeroDivisionError, list, dropwhile(isOdd, E(s)))
  652.  
  653.     def test_tee(self):
  654.         for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
  655.             for g in (G, I, Ig, S, L, R):
  656.                 it1, it2 = tee(g(s))
  657.                 self.assertEqual(list(it1), list(g(s)))
  658.                 self.assertEqual(list(it2), list(g(s)))
  659.             self.assertRaises(TypeError, tee, X(s))
  660.             self.assertRaises(TypeError, list, tee(N(s))[0])
  661.             self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
  662.  
  663. class LengthTransparency(unittest.TestCase):
  664.  
  665.     def test_repeat(self):
  666.         self.assertEqual(len(repeat(None, 50)), 50)
  667.         self.assertRaises(TypeError, len, repeat(None))
  668.  
  669. class RegressionTests(unittest.TestCase):
  670.  
  671.     def test_sf_793826(self):
  672.         # Fix Armin Rigo's successful efforts to wreak havoc
  673.  
  674.         def mutatingtuple(tuple1, f, tuple2):
  675.             # this builds a tuple t which is a copy of tuple1,
  676.             # then calls f(t), then mutates t to be equal to tuple2
  677.             # (needs len(tuple1) == len(tuple2)).
  678.             def g(value, first=[1]):
  679.                 if first:
  680.                     del first[:]
  681.                     f(z.next())
  682.                 return value
  683.             items = list(tuple2)
  684.             items[1:1] = list(tuple1)
  685.             gen = imap(g, items)
  686.             z = izip(*[gen]*len(tuple1))
  687.             z.next()
  688.  
  689.         def f(t):
  690.             global T
  691.             T = t
  692.             first[:] = list(T)
  693.  
  694.         first = []
  695.         mutatingtuple((1,2,3), f, (4,5,6))
  696.         second = list(T)
  697.         self.assertEqual(first, second)
  698.  
  699.  
  700.     def test_sf_950057(self):
  701.         # Make sure that chain() and cycle() catch exceptions immediately
  702.         # rather than when shifting between input sources
  703.  
  704.         def gen1():
  705.             hist.append(0)
  706.             yield 1
  707.             hist.append(1)
  708.             raise AssertionError
  709.             hist.append(2)
  710.  
  711.         def gen2(x):
  712.             hist.append(3)
  713.             yield 2
  714.             hist.append(4)
  715.             if x:
  716.                 raise StopIteration
  717.  
  718.         hist = []
  719.         self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
  720.         self.assertEqual(hist, [0,1])
  721.  
  722.         hist = []
  723.         self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
  724.         self.assertEqual(hist, [0,1])
  725.  
  726.         hist = []
  727.         self.assertRaises(AssertionError, list, cycle(gen1()))
  728.         self.assertEqual(hist, [0,1])
  729.  
  730. libreftest = """ Doctest for examples in the library reference: libitertools.tex
  731.  
  732.  
  733. >>> amounts = [120.15, 764.05, 823.14]
  734. >>> for checknum, amount in izip(count(1200), amounts):
  735. ...     print 'Check %d is for $%.2f' % (checknum, amount)
  736. ...
  737. Check 1200 is for $120.15
  738. Check 1201 is for $764.05
  739. Check 1202 is for $823.14
  740.  
  741. >>> import operator
  742. >>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
  743. ...    print cube
  744. ...
  745. 1
  746. 8
  747. 27
  748.  
  749. >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
  750. >>> for name in islice(reportlines, 3, None, 2):
  751. ...    print name.title()
  752. ...
  753. Alex
  754. Laura
  755. Martin
  756. Walter
  757. Samuele
  758.  
  759. >>> from operator import itemgetter
  760. >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
  761. >>> di = sorted(d.iteritems(), key=itemgetter(1))
  762. >>> for k, g in groupby(di, itemgetter(1)):
  763. ...     print k, map(itemgetter(0), g)
  764. ...
  765. 1 ['a', 'c', 'e']
  766. 2 ['b', 'd', 'f']
  767. 3 ['g']
  768.  
  769. # Find runs of consecutive numbers using groupby.  The key to the solution
  770. # is differencing with a range so that consecutive numbers all appear in
  771. # same group.
  772. >>> data = [ 1,  4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
  773. >>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
  774. ...     print map(operator.itemgetter(1), g)
  775. ...
  776. [1]
  777. [4, 5, 6]
  778. [10]
  779. [15, 16, 17, 18]
  780. [22]
  781. [25, 26, 27, 28]
  782.  
  783. >>> def take(n, seq):
  784. ...     return list(islice(seq, n))
  785.  
  786. >>> def enumerate(iterable):
  787. ...     return izip(count(), iterable)
  788.  
  789. >>> def tabulate(function):
  790. ...     "Return function(0), function(1), ..."
  791. ...     return imap(function, count())
  792.  
  793. >>> def iteritems(mapping):
  794. ...     return izip(mapping.iterkeys(), mapping.itervalues())
  795.  
  796. >>> def nth(iterable, n):
  797. ...     "Returns the nth item"
  798. ...     return list(islice(iterable, n, n+1))
  799.  
  800. >>> def all(seq, pred=bool):
  801. ...     "Returns True if pred(x) is True for every element in the iterable"
  802. ...     for elem in ifilterfalse(pred, seq):
  803. ...         return False
  804. ...     return True
  805.  
  806. >>> def any(seq, pred=bool):
  807. ...     "Returns True if pred(x) is True for at least one element in the iterable"
  808. ...     for elem in ifilter(pred, seq):
  809. ...         return True
  810. ...     return False
  811.  
  812. >>> def no(seq, pred=bool):
  813. ...     "Returns True if pred(x) is False for every element in the iterable"
  814. ...     for elem in ifilter(pred, seq):
  815. ...         return False
  816. ...     return True
  817.  
  818. >>> def quantify(seq, pred=bool):
  819. ...     "Count how many times the predicate is True in the sequence"
  820. ...     return sum(imap(pred, seq))
  821.  
  822. >>> def padnone(seq):
  823. ...     "Returns the sequence elements and then returns None indefinitely"
  824. ...     return chain(seq, repeat(None))
  825.  
  826. >>> def ncycles(seq, n):
  827. ...     "Returns the sequence elements n times"
  828. ...     return chain(*repeat(seq, n))
  829.  
  830. >>> def dotproduct(vec1, vec2):
  831. ...     return sum(imap(operator.mul, vec1, vec2))
  832.  
  833. >>> def flatten(listOfLists):
  834. ...     return list(chain(*listOfLists))
  835.  
  836. >>> def repeatfunc(func, times=None, *args):
  837. ...     "Repeat calls to func with specified arguments."
  838. ...     "   Example:  repeatfunc(random.random)"
  839. ...     if times is None:
  840. ...         return starmap(func, repeat(args))
  841. ...     else:
  842. ...         return starmap(func, repeat(args, times))
  843.  
  844. >>> def pairwise(iterable):
  845. ...     "s -> (s0,s1), (s1,s2), (s2, s3), ..."
  846. ...     a, b = tee(iterable)
  847. ...     try:
  848. ...         b.next()
  849. ...     except StopIteration:
  850. ...         pass
  851. ...     return izip(a, b)
  852.  
  853. This is not part of the examples but it tests to make sure the definitions
  854. perform as purported.
  855.  
  856. >>> take(10, count())
  857. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  858.  
  859. >>> list(enumerate('abc'))
  860. [(0, 'a'), (1, 'b'), (2, 'c')]
  861.  
  862. >>> list(islice(tabulate(lambda x: 2*x), 4))
  863. [0, 2, 4, 6]
  864.  
  865. >>> nth('abcde', 3)
  866. ['d']
  867.  
  868. >>> all([2, 4, 6, 8], lambda x: x%2==0)
  869. True
  870.  
  871. >>> all([2, 3, 6, 8], lambda x: x%2==0)
  872. False
  873.  
  874. >>> any([2, 4, 6, 8], lambda x: x%2==0)
  875. True
  876.  
  877. >>> any([1, 3, 5, 9], lambda x: x%2==0,)
  878. False
  879.  
  880. >>> no([1, 3, 5, 9], lambda x: x%2==0)
  881. True
  882.  
  883. >>> no([1, 2, 5, 9], lambda x: x%2==0)
  884. False
  885.  
  886. >>> quantify(xrange(99), lambda x: x%2==0)
  887. 50
  888.  
  889. >>> a = [[1, 2, 3], [4, 5, 6]]
  890. >>> flatten(a)
  891. [1, 2, 3, 4, 5, 6]
  892.  
  893. >>> list(repeatfunc(pow, 5, 2, 3))
  894. [8, 8, 8, 8, 8]
  895.  
  896. >>> import random
  897. >>> take(5, imap(int, repeatfunc(random.random)))
  898. [0, 0, 0, 0, 0]
  899.  
  900. >>> list(pairwise('abcd'))
  901. [('a', 'b'), ('b', 'c'), ('c', 'd')]
  902.  
  903. >>> list(pairwise([]))
  904. []
  905.  
  906. >>> list(pairwise('a'))
  907. []
  908.  
  909. >>> list(islice(padnone('abc'), 0, 6))
  910. ['a', 'b', 'c', None, None, None]
  911.  
  912. >>> list(ncycles('abc', 3))
  913. ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c']
  914.  
  915. >>> dotproduct([1,2,3], [4,5,6])
  916. 32
  917.  
  918. """
  919.  
  920. __test__ = {'libreftest' : libreftest}
  921.  
  922. def test_main(verbose=None):
  923.     test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
  924.                     RegressionTests, LengthTransparency)
  925.     test_support.run_unittest(*test_classes)
  926.  
  927.     # verify reference counting
  928.     if verbose and hasattr(sys, "gettotalrefcount"):
  929.         import gc
  930.         counts = [None] * 5
  931.         for i in xrange(len(counts)):
  932.             test_support.run_unittest(*test_classes)
  933.             gc.collect()
  934.             counts[i] = sys.gettotalrefcount()
  935.         print counts
  936.  
  937.     # doctest the examples in the library reference
  938.     test_support.run_doctest(sys.modules[__name__], verbose)
  939.  
  940. if __name__ == "__main__":
  941.     test_main(verbose=True)
  942.